24 Star 155 Fork 138

Ascend / MindSpeed

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
克隆/下载
rms_norm.md 679 Bytes
一键复制 编辑 原始数据 按行查看 历史
邓佳 提交于 2024-05-27 09:13 . !359change ascendspeed to mindspeed

rms_norm融合优化

问题分析

rms_norm常见于LLaMA、LLaMA2、Baichuan等LLM模型中用于归一化,由于torch侧没有提供rms_norm算子的接口,因此在模型中通常是以自定义的形式出现,这种形式的执行效率相对较低。

解决方法

MindSpeed对将rms_norm操作合并成一个算子,减少数据传输和临时存储。算子接口见link

使用场景

模型使用rms_norm作为归一化方式,脚本中设置了--normalization RMSNorm

使用方法

设置--used-fused-rmsnorm即可调用rms_norm融合算子。

使用效果

开启融合算子可以节省内存,提升性能。

1
https://gitee.com/ascend/MindSpeed.git
git@gitee.com:ascend/MindSpeed.git
ascend
MindSpeed
MindSpeed
master

搜索帮助